// Copyright 2025 The XLS Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef XLS_FUZZER_IR_FUZZER_IR_FUZZ_TEST_LIBRARY_H_
#define XLS_FUZZER_IR_FUZZER_IR_FUZZ_TEST_LIBRARY_H_

#include <cstdint>
#include <memory>
#include <ostream>
#include <string>
#include <string_view>
#include <vector>

#include "gtest/gtest.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "xls/fuzzer/ir_fuzzer/fuzz_program.pb.h"
#include "xls/ir/node.h"
#include "xls/ir/value.h"
#include "xls/passes/optimization_pass.h"

// Contains functions that are used in IR fuzz tests.

namespace xls {

// Used as a test name for Functions and Packages in GoogleTests.
inline constexpr std::string_view kFuzzTestName = "FuzzTest";

// Contains a Package which contains a randomly generated Function. The function
// was randomly generated by FuzzTest using the FuzzProgramProto protobuf
// instruction template.
struct FuzzPackage {
  std::unique_ptr<Package> p;
  FuzzProgramProto fuzz_program;
};

// Same as FuzzPackage but also can store arguments that are compatible with the
// function parameters.
struct FuzzPackageWithArgs {
  FuzzPackage fuzz_package;
  std::vector<std::vector<Value>> arg_sets;
};

void RecordFuzzInfo(const FuzzPackageWithArgs& args);
void RecordFuzzInfo(const FuzzPackage& args);
void RecordFuzzInfo(std::shared_ptr<Package> args);
void RecordFuzzInfo(const Package& args);

void OptimizationPassChangesOutputs(FuzzPackageWithArgs fuzz_package_with_args,
                                    OptimizationPass& pass);
absl::Status PassChangesOutputsWithProto(std::string proto_string,
                                         int64_t arg_set_count,
                                         OptimizationPass& pass);
absl::Status EquateProtoToIrTest(
    std::string proto_string, testing::Matcher<const Node*> expected_ir_node);
absl::StatusOr<FuzzPackage> BuildPackageFromProtoString(
    std::string proto_string);

FuzzPackageWithArgs GenArgSetsForPackage(FuzzPackage fuzz_package,
                                         int64_t arg_set_count);

// Intentionally override just this printer to avoid printing the protobuf.
void FuzzTestPrintSourceCode(const FuzzPackage& fp, std::ostream* os);
// Implement printing for arg sets.
void FuzzTestPrintSourceCode(const std::vector<std::vector<Value>>& fp,
                             std::ostream* os);

}  // namespace xls

#endif  // XLS_FUZZER_IR_FUZZER_IR_FUZZ_TEST_LIBRARY_H_
